#include "test.h"
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <malloc.h>
#include "bsl_err.h"
#include "bsl_sal.h"
#include "crypt_errno.h"
#include "crypt_eal_cipher.h"

static void PrintLastError(void) {
    const char *file = NULL;
    uint32_t line = 0;
    BSL_ERR_GetLastErrorFileLine(&file, &line); // 获取错误发生的文件名和行数
    printf("failed at file %s at line %d\r\n", file, (int)line);
}

void modes_sm4(CRYPT_CIPHER_AlgId cipher_id, CRYPT_PaddingType padding_id, bool padding) {
    uint8_t data[16] = {0};
    uint8_t iv[16] = {0};
    uint8_t key[16] = {0};
    uint32_t dataLen = sizeof(data);
    uint8_t cipherText[100];
    uint8_t plainText[100];
    uint32_t outTotalLen = 0;
    uint32_t outLen = sizeof(cipherText);
    uint32_t cipherTextLen;
    int32_t ret;

    printf("plain text to be encrypted: ");
    for (uint32_t i = 0; i < dataLen; i++) {
        printf("%02x", data[i]);
    }
    printf("\r\n");

    // Create a context.
    CRYPT_EAL_CipherCtx *ctx = CRYPT_EAL_CipherNewCtx(cipher_id);
    if (ctx == NULL) {
        PrintLastError();
        BSL_ERR_DeInit();
        return;
    }
    /*
     * During initialization, the last input parameter can be true or false. true indicates encryption,
     * and false indicates decryption.
     */
    ret = CRYPT_EAL_CipherInit(ctx, key, sizeof(key), iv, sizeof(iv), true);
    if (ret != CRYPT_SUCCESS) {
        // Output the error code. You can find the error information in **crypt_errno.h** based on the error code.
        printf("error code is %x\r\n", ret);
        PrintLastError();
        goto EXIT;
    }
    // Set the padding mode.
    if(padding) {
        ret = CRYPT_EAL_CipherSetPadding(ctx, padding_id);
        if (ret != CRYPT_SUCCESS) {
            printf("error code is %x\r\n", ret);
            PrintLastError();
            goto EXIT;
        }
    }
    /**
     * Enter the data to be calculated. This interface can be called for multiple times.
     * The input value of **outLen** is the length of the ciphertext,
     * and the output value is the amount of processed data.
     *
    */
    ret = CRYPT_EAL_CipherUpdate(ctx, data, dataLen, cipherText, &outLen);
    if (ret != CRYPT_SUCCESS) {
        printf("error code is %x\r\n", ret);
        PrintLastError();
        goto EXIT;
    }

    outTotalLen += outLen;
    outLen = sizeof(cipherText) - outTotalLen;

    ret = CRYPT_EAL_CipherFinal(ctx, cipherText + outTotalLen, &outLen);
    if (ret != CRYPT_SUCCESS) {
        printf("error code is %x\r\n", ret);
        PrintLastError();
        goto EXIT;
    }

    outTotalLen += outLen;
    printf("cipher text value is: ");

    for (uint32_t i = 0; i < outTotalLen; i++) {
        printf("%02x", cipherText[i]);
    }
    printf("\r\n");

    // Start decryption.
    cipherTextLen = outTotalLen;
    outTotalLen = 0;
    outLen = sizeof(plainText);

    // When initializing the decryption function, set the last input parameter to false.
    ret = CRYPT_EAL_CipherInit(ctx, key, sizeof(key), iv, sizeof(iv), false);
    if (ret != CRYPT_SUCCESS) {
        printf("error code is %x\r\n", ret);
        PrintLastError();
        goto EXIT;
    }

    // Set the padding mode, which must be the same as that for encryption.
    if(padding) {
        ret = CRYPT_EAL_CipherSetPadding(ctx, padding_id);
        if (ret != CRYPT_SUCCESS) {
            printf("error code is %x\r\n", ret);
            PrintLastError();
            goto EXIT;
        }
    }

    // Enter the ciphertext data.
    ret = CRYPT_EAL_CipherUpdate(ctx, cipherText, cipherTextLen, plainText, &outLen);
    if (ret != CRYPT_SUCCESS) {
        printf("error code is %x\r\n", ret);
        PrintLastError();
        goto EXIT;
    }
    outTotalLen += outLen;
    outLen = sizeof(plainText) - outTotalLen;

    // Decrypt the last segment of data and remove the filled content.
    ret = CRYPT_EAL_CipherFinal(ctx, plainText + outTotalLen, &outLen);
    if (ret != CRYPT_SUCCESS) {
        printf("error code is %x\r\n", ret);
        PrintLastError();
        goto EXIT;
    }

    outTotalLen += outLen;

    printf("decrypted plain text value is: ");
    for (uint32_t i = 0; i < outTotalLen; i++) {
        printf("%02x", plainText[i]);
    }
    printf("\r\n");

    if (outTotalLen != dataLen || memcmp(plainText, data, dataLen) != 0) {
        printf("plaintext comparison failed\r\n");
        goto EXIT;
    }
    printf("pass \r\n");

    EXIT:
    CRYPT_EAL_CipherFreeCtx(ctx);
    BSL_ERR_DeInit();
}

typedef struct {
    const uint8_t *data;
    uint32_t data_len;
    const uint8_t *key;
    uint32_t key_len;
    const uint8_t *iv;
    uint32_t iv_len;
    uint32_t buffer_size;
    int test_times;
} data_params_t;

typedef struct {
    CRYPT_CIPHER_AlgId cipher_id;
    CRYPT_PaddingType padding_id;
    bool use_padding;
    const char *mode_name;
} modes_params_t;

static void sm4_benchmark_core(const data_params_t *params, const  modes_params_t *modes_params) {
    printf("\r\n--- %s Mode Benchmark ---\r\n", modes_params->mode_name);

    uint8_t *cipherText = malloc(params->buffer_size);
    uint8_t *plainText = malloc(params->buffer_size);
    uint32_t outLen;
    uint32_t cipherTextLen = 0;
    int32_t ret;

    CRYPT_EAL_CipherCtx *ctx = CRYPT_EAL_CipherNewCtx(modes_params->cipher_id);
    if (ctx) {
        // 先执行一次加密获取实际密文长度
        CRYPT_EAL_CipherInit(ctx, params->key, params->key_len, params->iv, params->iv_len, true);
        if (modes_params->use_padding) {
            CRYPT_EAL_CipherSetPadding(ctx, modes_params->padding_id);
        }

        uint32_t tempOutLen = params->buffer_size;
        ret = CRYPT_EAL_CipherUpdate(ctx, params->data, params->data_len, cipherText, &tempOutLen);
        if (ret == CRYPT_SUCCESS) {
            cipherTextLen = tempOutLen;
            tempOutLen = params->buffer_size - cipherTextLen;
            ret = CRYPT_EAL_CipherFinal(ctx, cipherText + cipherTextLen, &tempOutLen);
            if (ret == CRYPT_SUCCESS) {
                cipherTextLen += tempOutLen;
            }
        }

//        printf("%s encrypted length: %lu bytes\r\n", modes_params->mode_name, cipherTextLen);

        printf("%-25s, %15s, %20s, %20s\r\n", "algorithm operation", "time elapsed(s)", "run times", "ops");
        // 加密性能测试
        BENCH_TIMES_VOID({
                             CRYPT_EAL_CipherInit(ctx, params->key, params->key_len, params->iv, params->iv_len, true);
                             if (modes_params->use_padding) {
                                 CRYPT_EAL_CipherSetPadding(ctx, modes_params->padding_id);
                             }
                             outLen = params->buffer_size;
                             CRYPT_EAL_CipherUpdate(ctx, params->data, params->data_len, cipherText, &outLen);
                             CRYPT_EAL_CipherFinal(ctx, cipherText + outLen, &outLen);
                         }, params->test_times, "Encrypt");

        // 解密性能测试
        BENCH_TIMES_VOID({
                             CRYPT_EAL_CipherInit(ctx, params->key, params->key_len, params->iv, params->iv_len, false);
                             if (modes_params->use_padding) {
                                 CRYPT_EAL_CipherSetPadding(ctx, modes_params->padding_id);
                             }
                             outLen = params->buffer_size;
                             CRYPT_EAL_CipherUpdate(ctx, cipherText, cipherTextLen, plainText, &outLen);
                             CRYPT_EAL_CipherFinal(ctx, plainText + outLen, &outLen);
                         }, params->test_times, "Decrypt");

        CRYPT_EAL_CipherFreeCtx(ctx);
    }

    free(cipherText);
    free(plainText);
}

// 原来的三个函数现在简化为对统一函数的调用
static void CBC_benchmark(const data_params_t *params) {
     modes_params_t modes_params = {
            .cipher_id = CRYPT_CIPHER_SM4_CBC,
            .padding_id = CRYPT_PADDING_PKCS7,
            .use_padding = true,
            .mode_name = "CBC"
    };
    sm4_benchmark_core(params, &modes_params);
}

static void ECB_benchmark(const data_params_t *params) {
     modes_params_t modes_params = {
            .cipher_id = CRYPT_CIPHER_SM4_ECB,
            .padding_id = CRYPT_PADDING_PKCS7,
            .use_padding = true,
            .mode_name = "ECB"
    };
    sm4_benchmark_core(params, &modes_params);
}

static void CTR_benchmark(const data_params_t *params) {
     modes_params_t modes_params = {
            .cipher_id = CRYPT_CIPHER_SM4_CTR,
            .padding_id = CRYPT_PADDING_PKCS7, // CTR模式通常不需要填充，但为了统一接口保留
            .use_padding = false,              // CTR模式不设置填充
            .mode_name = "CTR"
    };
    sm4_benchmark_core(params, &modes_params);
}

/**
 * @brief SM4加密算法性能测试主接口
 * @param data_size 测试数据大小（字节）
 * @param test_times 每个测试的迭代次数
 */
static void benchmark_sm4(uint32_t data_size, int test_times) {
    printf("\r\n=== SM4 Performance Benchmark (%lu bytes, %d iterations) ===\r\n",
           data_size, test_times);

    // 分配和初始化测试数据
    uint8_t *data = malloc(data_size);
    uint8_t iv[16] = {0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
                      0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f};
    uint8_t key[16] = {0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef,
                       0xfe, 0xdc, 0xba, 0x98, 0x76, 0x54, 0x32, 0x10};

    // 填充测试数据
    for (uint32_t i = 0; i < data_size; i++) {
        data[i] = (i & 0xFF);
    }

    // 设置性能测试参数
    data_params_t params = {
            .data = data,
            .data_len = data_size,
            .key = key,
            .key_len = sizeof(key),
            .iv = iv,
            .iv_len = sizeof(iv),
            .buffer_size = data_size + 32,
            .test_times = test_times
    };

    printf("Benchmark parameters:\r\n");
    printf("  Data length: %lu bytes\r\n", params.data_len);
    printf("  Key length: %lu bytes\r\n", params.key_len);
    printf("  IV length: %lu bytes\r\n", params.iv_len);
    printf("  Buffer size: %lu bytes\r\n", params.buffer_size);
    printf("  Test iterations: %d\r\n", params.test_times);

    // 执行各模式性能测试
    CBC_benchmark(&params);
    ECB_benchmark(&params);
    CTR_benchmark(&params);

    // 清理资源
    free(data);

    printf("\r\n=== Benchmark Complete ===\r\n");
}

void test_sm4(void) {
//    modes_sm4(CRYPT_CIPHER_SM4_CBC, CRYPT_PADDING_PKCS7, true);
//    modes_sm4(CRYPT_CIPHER_SM4_ECB, CRYPT_PADDING_PKCS7, true);
//    modes_sm4(CRYPT_CIPHER_SM4_CTR, CRYPT_PADDING_NONE, false);

//    benchmark_sm4(32, 1000);
    benchmark_sm4(4096, 100);
}